import matplotlib.pyplot as plt
import numpy as np

# 创建相关性矩阵
np.random.seed(42)
data = np.random.rand(5, 5)

fig, ax = plt.subplots(figsize=(8, 6))
im = ax.imshow(data, cmap='viridis')

# 添加颜色条
cbar = ax.figure.colorbar(im, ax=ax)
cbar.ax.set_ylabel('数值', rotation=-90, va="bottom")

# 设置刻度标签
ax.set_xticks(np.arange(5))
ax.set_yticks(np.arange(5))
ax.set_xticklabels(['A', 'B', 'C', 'D', 'E'])
ax.set_yticklabels(['1', '2', '3', '4', '5'])

# 在热力图中显示数值
for i in range(5):
    for j in range(5):
        text = ax.text(j, i, f'{data[i, j]:.2f}',
                       ha="center", va="center", color="w")

ax.set_title("热力图示例")
plt.show()